{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "67c8afd7",
      "metadata": {
        "id": "67c8afd7"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/docs/examples/agent/multi_document_agents.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "43497beb-817d-4366-9156-f4d7f0d44942",
      "metadata": {
        "id": "43497beb-817d-4366-9156-f4d7f0d44942"
      },
      "source": [
        "# Multi-Document Agents\n",
        "\n",
        "In this guide, you learn towards setting up an agent that can effectively answer different types of questions over a larger set of documents.\n",
        "\n",
        "These questions include the following\n",
        "\n",
        "- QA over a specific doc\n",
        "- QA comparing different docs\n",
        "- Summaries over a specific doc\n",
        "- Comparing summaries between different docs\n",
        "\n",
        "We do this with the following architecture:\n",
        "\n",
        "- setup a \"document agent\" over each Document: each doc agent can do QA/summarization within its doc\n",
        "- setup a top-level agent over this set of document agents. Do tool retrieval and then do CoT over the set of tools to answer a question."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9be00aba-b6c5-4940-9825-81c5d2cd2f0b",
      "metadata": {
        "id": "9be00aba-b6c5-4940-9825-81c5d2cd2f0b"
      },
      "source": [
        "## Setup and Download Data\n",
        "\n",
        "In this section, we'll define imports and then download Wikipedia articles about different cities. Each article is stored separately.\n",
        "\n",
        "We load in 18 cities - this is not quite at the level of \"hundreds\" of documents but its still large enough to warrant some top-level document retrieval!"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5d81f93c",
      "metadata": {
        "id": "5d81f93c"
      },
      "source": [
        "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4110483a",
      "metadata": {
        "id": "4110483a"
      },
      "outputs": [],
      "source": [
        "%pip install llama-index-agent-openai\n",
        "%pip install llama-index-embeddings-openai\n",
        "%pip install llama-index-llms-openai"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f0fb1340",
      "metadata": {
        "id": "f0fb1340"
      },
      "outputs": [],
      "source": [
        "!pip install llama-index"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e41e9905-77a9-44c5-88ac-c7a4d08a4612",
      "metadata": {
        "id": "e41e9905-77a9-44c5-88ac-c7a4d08a4612"
      },
      "outputs": [],
      "source": [
        "from llama_index.core import (\n",
        "    VectorStoreIndex,\n",
        "    SimpleKeywordTableIndex,\n",
        "    SimpleDirectoryReader,\n",
        ")\n",
        "from llama_index.core import SummaryIndex\n",
        "from llama_index.core.schema import IndexNode\n",
        "from llama_index.core.tools import QueryEngineTool, ToolMetadata\n",
        "from llama_index.llms.openai import OpenAI\n",
        "from llama_index.core.callbacks import CallbackManager"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e4343cf7-eec9-4a67-b5be-c72dbe3280a7",
      "metadata": {
        "id": "e4343cf7-eec9-4a67-b5be-c72dbe3280a7"
      },
      "outputs": [],
      "source": [
        "wiki_titles = [\n",
        "    \"Toronto\",\n",
        "    \"Seattle\",\n",
        "    \"Chicago\",\n",
        "    \"Boston\",\n",
        "    \"Houston\",\n",
        "    \"Tokyo\",\n",
        "    \"Berlin\",\n",
        "    \"Lisbon\",\n",
        "    \"Paris\",\n",
        "    \"London\",\n",
        "    \"Atlanta\",\n",
        "    \"Munich\",\n",
        "    \"Shanghai\",\n",
        "    \"Beijing\",\n",
        "    \"Copenhagen\",\n",
        "    \"Moscow\",\n",
        "    \"Cairo\",\n",
        "    \"Karachi\",\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6d261882-6793-4eca-ad93-d94d2061e388",
      "metadata": {
        "id": "6d261882-6793-4eca-ad93-d94d2061e388"
      },
      "outputs": [],
      "source": [
        "from pathlib import Path\n",
        "\n",
        "import requests\n",
        "\n",
        "for title in wiki_titles:\n",
        "    response = requests.get(\n",
        "        \"https://en.wikipedia.org/w/api.php\",\n",
        "        params={\n",
        "            \"action\": \"query\",\n",
        "            \"format\": \"json\",\n",
        "            \"titles\": title,\n",
        "            \"prop\": \"extracts\",\n",
        "            # 'exintro': True,\n",
        "            \"explaintext\": True,\n",
        "        },\n",
        "    ).json()\n",
        "    page = next(iter(response[\"query\"][\"pages\"].values()))\n",
        "    wiki_text = page[\"extract\"]\n",
        "\n",
        "    data_path = Path(\"data\")\n",
        "    if not data_path.exists():\n",
        "        Path.mkdir(data_path)\n",
        "\n",
        "    with open(data_path / f\"{title}.txt\", \"w\") as fp:\n",
        "        fp.write(wiki_text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5bf0c13b-0d77-43e8-8c1c-84258299a494",
      "metadata": {
        "id": "5bf0c13b-0d77-43e8-8c1c-84258299a494"
      },
      "outputs": [],
      "source": [
        "# Load all wiki documents\n",
        "city_docs = {}\n",
        "for wiki_title in wiki_titles:\n",
        "    city_docs[wiki_title] = SimpleDirectoryReader(\n",
        "        input_files=[f\"data/{wiki_title}.txt\"]\n",
        "    ).load_data()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6189aaf4-2eb7-40bc-9e83-79ce4f221b4b",
      "metadata": {
        "id": "6189aaf4-2eb7-40bc-9e83-79ce4f221b4b"
      },
      "source": [
        "Define Global LLM and Embeddings"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "62ede51d",
      "metadata": {
        "id": "62ede51d"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "dd6e5e48-91b9-4701-a85d-d98c92323350",
      "metadata": {
        "id": "dd6e5e48-91b9-4701-a85d-d98c92323350"
      },
      "outputs": [],
      "source": [
        "from llama_index.llms.openai import OpenAI\n",
        "from llama_index.embeddings.openai import OpenAIEmbedding\n",
        "from llama_index.core import Settings\n",
        "\n",
        "Settings.llm = OpenAI(temperature=0, model=\"gpt-3.5-turbo\")\n",
        "Settings.embed_model = OpenAIEmbedding(model=\"text-embedding-ada-002\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4eeef31a-fc25-4367-a5ba-945f81d04cf9",
      "metadata": {
        "id": "4eeef31a-fc25-4367-a5ba-945f81d04cf9"
      },
      "source": [
        "## Building Multi-Document Agents\n",
        "\n",
        "In this section we show you how to construct the multi-document agent. We first build a document agent for each document, and then define the top-level parent agent with an object index."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "976cd798-2e8d-474c-922a-51b12c5c6f36",
      "metadata": {
        "id": "976cd798-2e8d-474c-922a-51b12c5c6f36"
      },
      "source": [
        "### Build Document Agent for each Document\n",
        "\n",
        "In this section we define \"document agents\" for each document.\n",
        "\n",
        "We define both a vector index (for semantic search) and summary index (for summarization) for each document. The two query engines are then converted into tools that are passed to an OpenAI function calling agent.\n",
        "\n",
        "This document agent can dynamically choose to perform semantic search or summarization within a given document.\n",
        "\n",
        "We create a separate document agent for each city."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "eacdf3a7-cfe3-4c2b-9037-b28a065ed148",
      "metadata": {
        "id": "eacdf3a7-cfe3-4c2b-9037-b28a065ed148"
      },
      "outputs": [],
      "source": [
        "from llama_index.agent.openai import OpenAIAgent\n",
        "from llama_index.core import load_index_from_storage, StorageContext\n",
        "from llama_index.core.node_parser import SentenceSplitter\n",
        "import os\n",
        "\n",
        "node_parser = SentenceSplitter()\n",
        "\n",
        "# Build agents dictionary\n",
        "agents = {}\n",
        "query_engines = {}\n",
        "\n",
        "# this is for the baseline\n",
        "all_nodes = []\n",
        "\n",
        "for idx, wiki_title in enumerate(wiki_titles):\n",
        "    nodes = node_parser.get_nodes_from_documents(city_docs[wiki_title])\n",
        "    all_nodes.extend(nodes)\n",
        "\n",
        "    if not os.path.exists(f\"./data/{wiki_title}\"):\n",
        "        # build vector index\n",
        "        vector_index = VectorStoreIndex(nodes)\n",
        "        vector_index.storage_context.persist(\n",
        "            persist_dir=f\"./data/{wiki_title}\"\n",
        "        )\n",
        "    else:\n",
        "        vector_index = load_index_from_storage(\n",
        "            StorageContext.from_defaults(persist_dir=f\"./data/{wiki_title}\"),\n",
        "        )\n",
        "\n",
        "    # build summary index\n",
        "    summary_index = SummaryIndex(nodes)\n",
        "    # define query engines\n",
        "    vector_query_engine = vector_index.as_query_engine(llm=Settings.llm)\n",
        "    summary_query_engine = summary_index.as_query_engine(llm=Settings.llm)\n",
        "\n",
        "    # define tools\n",
        "    query_engine_tools = [\n",
        "        QueryEngineTool(\n",
        "            query_engine=vector_query_engine,\n",
        "            metadata=ToolMetadata(\n",
        "                name=\"vector_tool\",\n",
        "                description=(\n",
        "                    \"Useful for questions related to specific aspects of\"\n",
        "                    f\" {wiki_title} (e.g. the history, arts and culture,\"\n",
        "                    \" sports, demographics, or more).\"\n",
        "                ),\n",
        "            ),\n",
        "        ),\n",
        "        QueryEngineTool(\n",
        "            query_engine=summary_query_engine,\n",
        "            metadata=ToolMetadata(\n",
        "                name=\"summary_tool\",\n",
        "                description=(\n",
        "                    \"Useful for any requests that require a holistic summary\"\n",
        "                    f\" of EVERYTHING about {wiki_title}. For questions about\"\n",
        "                    \" more specific sections, please use the vector_tool.\"\n",
        "                ),\n",
        "            ),\n",
        "        ),\n",
        "    ]\n",
        "\n",
        "    # build agent\n",
        "    function_llm = OpenAI(model=\"gpt-4\")\n",
        "    agent = OpenAIAgent.from_tools(\n",
        "        query_engine_tools,\n",
        "        llm=function_llm,\n",
        "        verbose=True,\n",
        "        system_prompt=f\"\"\"\\\n",
        "You are a specialized agent designed to answer queries about {wiki_title}.\n",
        "You must ALWAYS use at least one of the tools provided when answering a question; do NOT rely on prior knowledge.\\\n",
        "\"\"\",\n",
        "    )\n",
        "\n",
        "    agents[wiki_title] = agent\n",
        "    query_engines[wiki_title] = vector_index.as_query_engine(\n",
        "        similarity_top_k=2\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "899ca55b-0c02-429b-a765-8e4f806d503f",
      "metadata": {
        "id": "899ca55b-0c02-429b-a765-8e4f806d503f"
      },
      "source": [
        "### Build Retriever-Enabled OpenAI Agent\n",
        "\n",
        "We build a top-level agent that can orchestrate across the different document agents to answer any user query.\n",
        "\n",
        "This agent takes in all document agents as tools. This specific agent `RetrieverOpenAIAgent` performs tool retrieval before tool use (unlike a default agent that tries to put all tools in the prompt).\n",
        "\n",
        "Here we use a top-k retriever, but we encourage you to customize the tool retriever method!\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6884ff15-bf40-4bdd-a1e3-58cbd056a12a",
      "metadata": {
        "id": "6884ff15-bf40-4bdd-a1e3-58cbd056a12a"
      },
      "outputs": [],
      "source": [
        "# define tool for each document agent\n",
        "all_tools = []\n",
        "for wiki_title in wiki_titles:\n",
        "    wiki_summary = (\n",
        "        f\"This content contains Wikipedia articles about {wiki_title}. Use\"\n",
        "        f\" this tool if you want to answer any questions about {wiki_title}.\\n\"\n",
        "    )\n",
        "    doc_tool = QueryEngineTool(\n",
        "        query_engine=agents[wiki_title],\n",
        "        metadata=ToolMetadata(\n",
        "            name=f\"tool_{wiki_title}\",\n",
        "            description=wiki_summary,\n",
        "        ),\n",
        "    )\n",
        "    all_tools.append(doc_tool)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "da80793c-54b9-43e2-b34a-27412156268a",
      "metadata": {
        "id": "da80793c-54b9-43e2-b34a-27412156268a"
      },
      "outputs": [],
      "source": [
        "# define an \"object\" index and retriever over these tools\n",
        "from llama_index.core import VectorStoreIndex\n",
        "from llama_index.core.objects import ObjectIndex\n",
        "\n",
        "obj_index = ObjectIndex.from_objects(\n",
        "    all_tools,\n",
        "    index_cls=VectorStoreIndex,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fed38942-1e37-4c61-89fa-d2ef41151831",
      "metadata": {
        "id": "fed38942-1e37-4c61-89fa-d2ef41151831"
      },
      "outputs": [],
      "source": [
        "from llama_index.agent.openai import OpenAIAgent\n",
        "\n",
        "top_agent = OpenAIAgent.from_tools(\n",
        "    tool_retriever=obj_index.as_retriever(similarity_top_k=3),\n",
        "    system_prompt=\"\"\" \\\n",
        "You are an agent designed to answer queries about a set of given cities.\n",
        "Please always use the tools provided to answer a question. Do not rely on prior knowledge.\\\n",
        "\n",
        "\"\"\",\n",
        "    verbose=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "aa32b97c-6779-4b60-823d-6ca3be6f358a",
      "metadata": {
        "id": "aa32b97c-6779-4b60-823d-6ca3be6f358a"
      },
      "source": [
        "### Define Baseline Vector Store Index\n",
        "\n",
        "As a point of comparison, we define a \"naive\" RAG pipeline which dumps all docs into a single vector index collection.\n",
        "\n",
        "We set the top_k = 4"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "78e7b48c-687a-4585-a730-e651212de3b5",
      "metadata": {
        "id": "78e7b48c-687a-4585-a730-e651212de3b5"
      },
      "outputs": [],
      "source": [
        "base_index = VectorStoreIndex(all_nodes)\n",
        "base_query_engine = base_index.as_query_engine(similarity_top_k=4)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8dedb927-a992-4f21-a0fb-4ce4361adcb3",
      "metadata": {
        "id": "8dedb927-a992-4f21-a0fb-4ce4361adcb3"
      },
      "source": [
        "## Running Example Queries\n",
        "\n",
        "Let's run some example queries, ranging from QA / summaries over a single document to QA / summarization over multiple documents."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8e743c62-7dd8-4ac9-85a5-f1cbc112a79c",
      "metadata": {
        "id": "8e743c62-7dd8-4ac9-85a5-f1cbc112a79c",
        "outputId": "4917bf87-497b-4746-fe75-bf15e2a96d66"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "=== Calling Function ===\n",
            "Calling function: tool_Boston with args: {\n",
            "  \"input\": \"arts and culture\"\n",
            "}\n",
            "=== Calling Function ===\n",
            "Calling function: vector_tool with args: {\n",
            "  \"input\": \"arts and culture\"\n",
            "}\n",
            "Got output: Boston is known for its vibrant arts and culture scene. The city is home to a number of performing arts organizations, including the Boston Ballet, Boston Lyric Opera Company, Opera Boston, Boston Baroque, and the Handel and Haydn Society. There are also several theaters in or near the Theater District, such as the Cutler Majestic Theatre, Citi Performing Arts Center, the Colonial Theater, and the Orpheum Theatre. Boston is a center for contemporary classical music, with groups like the Boston Modern Orchestra Project and Boston Musica Viva. The city also hosts major annual events, such as First Night, the Boston Early Music Festival, and the Boston Arts Festival. In addition, Boston has several art museums and galleries, including the Museum of Fine Arts, the Isabella Stewart Gardner Museum, and the Institute of Contemporary Art.\n",
            "========================\n",
            "Got output: Boston is renowned for its vibrant arts and culture scene. It is home to numerous performing arts organizations, including the Boston Ballet, Boston Lyric Opera Company, Opera Boston, Boston Baroque, and the Handel and Haydn Society. The city's Theater District houses several theaters, such as the Cutler Majestic Theatre, Citi Performing Arts Center, the Colonial Theater, and the Orpheum Theatre.\n",
            "\n",
            "Boston is also a hub for contemporary classical music, with groups like the Boston Modern Orchestra Project and Boston Musica Viva. The city hosts major annual events, such as First Night, the Boston Early Music Festival, and the Boston Arts Festival, which contribute to its cultural richness.\n",
            "\n",
            "In terms of visual arts, Boston boasts several art museums and galleries. The Museum of Fine Arts, the Isabella Stewart Gardner Museum, and the Institute of Contemporary Art are among the most notable. These institutions offer a wide range of art collections, from ancient to contemporary, attracting art enthusiasts from around the world.\n",
            "========================\n"
          ]
        }
      ],
      "source": [
        "# should use Boston agent -> vector tool\n",
        "response = top_agent.query(\"Tell me about the arts and culture in Boston\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a4ce2a76-5779-4acf-9337-69109dae7fd6",
      "metadata": {
        "id": "a4ce2a76-5779-4acf-9337-69109dae7fd6",
        "outputId": "0fd225d7-4ea0-4675-84af-5a192ab9db89"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Boston has a rich arts and culture scene, with a variety of performing arts organizations and venues. The city is home to renowned institutions such as the Boston Ballet, Boston Lyric Opera Company, Opera Boston, Boston Baroque, and the Handel and Haydn Society. The Theater District in Boston is a hub for theatrical performances, with theaters like the Cutler Majestic Theatre, Citi Performing Arts Center, Colonial Theater, and Orpheum Theatre.\n",
            "\n",
            "In addition to performing arts, Boston also has a thriving contemporary classical music scene, with groups like the Boston Modern Orchestra Project and Boston Musica Viva. The city hosts several annual events that celebrate the arts, including First Night, the Boston Early Music Festival, and the Boston Arts Festival.\n",
            "\n",
            "Boston is also known for its visual arts scene, with a number of art museums and galleries. The Museum of Fine Arts, the Isabella Stewart Gardner Museum, and the Institute of Contemporary Art are among the notable institutions in the city. These museums offer a diverse range of art collections, spanning from ancient to contemporary art, and attract art enthusiasts from around the world.\n"
          ]
        }
      ],
      "source": [
        "print(response)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "af28b422-fb73-4b59-9e77-3ba3afa87795",
      "metadata": {
        "id": "af28b422-fb73-4b59-9e77-3ba3afa87795",
        "outputId": "de5dc8ed-b802-4d94-9d4d-75df8c1907fb"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Boston has a rich arts and culture scene. The city is home to a variety of performing arts organizations, such as the Boston Ballet, Boston Lyric Opera Company, Opera Boston, Boston Baroque, and the Handel and Haydn Society. Additionally, there are numerous contemporary classical music groups associated with the city's conservatories and universities, like the Boston Modern Orchestra Project and Boston Musica Viva. The Theater District in Boston is a hub for theater, with notable venues including the Cutler Majestic Theatre, Citi Performing Arts Center, the Colonial Theater, and the Orpheum Theatre. Boston also hosts several significant annual events, including First Night, the Boston Early Music Festival, the Boston Arts Festival, and the Boston gay pride parade and festival. The city is renowned for its historic sites connected to the American Revolution, as well as its art museums and galleries, such as the Museum of Fine Arts, Isabella Stewart Gardner Museum, and the Institute of Contemporary Art.\n"
          ]
        }
      ],
      "source": [
        "# baseline\n",
        "response = base_query_engine.query(\n",
        "    \"Tell me about the arts and culture in Boston\"\n",
        ")\n",
        "print(str(response))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "aa3d98ab-cb82-4473-ab2b-bc8a17e1b86a",
      "metadata": {
        "id": "aa3d98ab-cb82-4473-ab2b-bc8a17e1b86a",
        "outputId": "5e7f7c3a-0fa6-45f5-d0b3-a5cf105eaa44"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "=== Calling Function ===\n",
            "Calling function: tool_Houston with args: {\n",
            "  \"input\": \"positive aspects\"\n",
            "}\n",
            "=== Calling Function ===\n",
            "Calling function: summary_tool with args: {\n",
            "  \"input\": \"positive aspects\"\n",
            "}\n",
            "Got output: Houston has many positive aspects that make it an attractive place to live and visit. The city's diverse population, with people from different ethnic and religious backgrounds, adds to its cultural richness and inclusiveness. Additionally, Houston is home to the Texas Medical Center, which is the largest concentration of healthcare and research institutions in the world. The presence of NASA's Johnson Space Center also highlights Houston's importance in the fields of medicine and space exploration. The city's strong economy, supported by industries such as energy, manufacturing, aeronautics, and transportation, provides numerous economic opportunities for residents and visitors alike. Furthermore, Houston has a thriving visual and performing arts scene, including a theater district and a variety of museums and galleries. Overall, Houston's diverse community, cultural attractions, and economic prospects make it an exceptionally appealing city.\n",
            "========================\n",
            "Got output: Houston has numerous positive aspects that make it a desirable place to live and visit. Some of these include:\n",
            "\n",
            "1. **Diversity**: Houston is known for its diverse population, with people from different ethnic and religious backgrounds. This diversity adds to the city's cultural richness and inclusiveness.\n",
            "\n",
            "2. **Healthcare and Research Institutions**: The city is home to the Texas Medical Center, the largest concentration of healthcare and research institutions in the world. This makes Houston a hub for medical innovation and healthcare services.\n",
            "\n",
            "3. **Space Exploration**: Houston is also known for NASA's Johnson Space Center, highlighting the city's significant role in space exploration.\n",
            "\n",
            "4. **Strong Economy**: Houston's economy is robust and diverse, supported by industries such as energy, manufacturing, aeronautics, and transportation. This provides numerous economic opportunities for its residents.\n",
            "\n",
            "5. **Arts and Culture**: The city has a thriving visual and performing arts scene, with a theater district and a variety of museums and galleries. This makes Houston a vibrant place for art lovers and creatives.\n",
            "\n",
            "Overall, these aspects contribute to making Houston an appealing and dynamic city.\n",
            "========================\n"
          ]
        }
      ],
      "source": [
        "# should use Houston agent -> vector tool\n",
        "response = top_agent.query(\n",
        "    \"Give me a summary of all the positive aspects of Houston\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d476c54b-98af-4d2a-8f17-4baa37d0d360",
      "metadata": {
        "id": "d476c54b-98af-4d2a-8f17-4baa37d0d360",
        "outputId": "205d3982-7d64-4878-9072-97bbd954e4d7"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Houston has numerous positive aspects that make it a desirable place to live and visit. Some of these include:\n",
            "\n",
            "1. Diversity: Houston is known for its diverse population, with people from different ethnic and religious backgrounds. This diversity adds to the city's cultural richness and inclusiveness.\n",
            "\n",
            "2. Healthcare and Research Institutions: The city is home to the Texas Medical Center, the largest concentration of healthcare and research institutions in the world. This makes Houston a hub for medical innovation and healthcare services.\n",
            "\n",
            "3. Space Exploration: Houston is also known for NASA's Johnson Space Center, highlighting the city's significant role in space exploration.\n",
            "\n",
            "4. Strong Economy: Houston's economy is robust and diverse, supported by industries such as energy, manufacturing, aeronautics, and transportation. This provides numerous economic opportunities for its residents.\n",
            "\n",
            "5. Arts and Culture: The city has a thriving visual and performing arts scene, with a theater district and a variety of museums and galleries. This makes Houston a vibrant place for art lovers and creatives.\n",
            "\n",
            "Overall, these aspects contribute to making Houston an appealing and dynamic city.\n"
          ]
        }
      ],
      "source": [
        "print(response)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c89b3581-ff70-4069-8cbb-f2791756128a",
      "metadata": {
        "id": "c89b3581-ff70-4069-8cbb-f2791756128a",
        "outputId": "e35d3494-d6da-4242-f6a1-9258c37a715c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Houston has several positive aspects that contribute to its reputation as a thriving city. It is home to a diverse and growing international community, with a large number of foreign banks and consular offices representing 92 countries. The city has received numerous accolades, including being ranked as one of the best cities for employment, college graduates, and homebuyers. Houston has a strong economy, with a broad industrial base in sectors such as energy, manufacturing, aeronautics, and healthcare. It is also a major center for the oil and gas industry and has the second-most Fortune 500 headquarters in the United States. The city's cultural scene is vibrant, with a variety of annual events celebrating different cultures, as well as a reputation for diverse and excellent food. Houston is known for its world-class museums and performing arts scene. Additionally, the city has made significant investments in renewable energy sources like wind and solar. Overall, Houston offers a high quality of life, reasonable living costs, and abundant employment opportunities.\n"
          ]
        }
      ],
      "source": [
        "# baseline\n",
        "response = base_query_engine.query(\n",
        "    \"Give me a summary of all the positive aspects of Houston\"\n",
        ")\n",
        "print(str(response))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c9d7dcec-5d1a-485a-b40f-aa9604fd92df",
      "metadata": {
        "id": "c9d7dcec-5d1a-485a-b40f-aa9604fd92df"
      },
      "outputs": [],
      "source": [
        "# baseline: the response doesn't quite match the sources...\n",
        "response.source_nodes[1].get_content()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ee6ef20c-3ccc-46c3-ad87-667138d78d5d",
      "metadata": {
        "id": "ee6ef20c-3ccc-46c3-ad87-667138d78d5d",
        "outputId": "d1f79d69-dcc0-428b-ed79-b238c61fc6d1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "=== Calling Function ===\n",
            "Calling function: tool_Houston with args: {\n",
            "  \"input\": \"demographics\"\n",
            "}\n",
            "=== Calling Function ===\n",
            "Calling function: vector_tool with args: {\n",
            "  \"input\": \"demographics\"\n",
            "}\n",
            "Got output: Houston is a majority-minority city with a diverse population. According to the U.S. Census Bureau, in 2019, non-Hispanic whites made up 23.3% of the population, Hispanics and Latino Americans 45.8%, Blacks or African Americans 22.4%, and Asian Americans 6.5%. The largest Hispanic or Latino American ethnic group in the city is Mexican Americans, followed by Puerto Ricans and Cuban Americans. Houston is also home to the largest African American community west of the Mississippi River. Additionally, Houston has a growing Muslim population, with Muslims estimated to make up 1.2% of the city's population. The city is known for its LGBT community and is home to one of the largest pride parades in the United States. The Hindu, Sikh, and Buddhist communities are also growing in Houston. Overall, Houston is considered one of the most ethnically and culturally diverse metropolitan areas in the country.\n",
            "========================\n",
            "Got output: Houston is a majority-minority city with a diverse population. According to the U.S. Census Bureau, in 2019, non-Hispanic whites made up 23.3% of the population, Hispanics and Latino Americans 45.8%, Blacks or African Americans 22.4%, and Asian Americans 6.5%. The largest Hispanic or Latino American ethnic group in the city is Mexican Americans, followed by Puerto Ricans and Cuban Americans. \n",
            "\n",
            "Houston is also home to the largest African American community west of the Mississippi River. Additionally, Houston has a growing Muslim population, with Muslims estimated to make up 1.2% of the city's population. The city is known for its LGBT community and is home to one of the largest pride parades in the United States. The Hindu, Sikh, and Buddhist communities are also growing in Houston. \n",
            "\n",
            "Overall, Houston is considered one of the most ethnically and culturally diverse metropolitan areas in the country.\n",
            "========================\n",
            "=== Calling Function ===\n",
            "Calling function: tool_Chicago with args: {\n",
            "  \"input\": \"demographics\"\n",
            "}\n",
            "=== Calling Function ===\n",
            "Calling function: vector_tool with args: {\n",
            "  \"input\": \"demographics\"\n",
            "}\n",
            "Got output: Chicago has a diverse demographic makeup. It experienced rapid population growth during its early years, becoming one of the fastest-growing cities in the world. Waves of immigrants from various European countries, as well as African Americans from the American South, contributed to the city's population growth. Over time, Chicago's population has fluctuated, with a decline in the latter half of the 20th century followed by a rise in recent years. As of the latest census estimates, the largest racial or ethnic groups in Chicago are non-Hispanic White, Black, and Hispanic. Additionally, Chicago has a significant LGBT population and is known for its cultural diversity.\n",
            "========================\n",
            "Got output: Chicago is known for its diverse demographic makeup. The city experienced rapid population growth during its early years, with immigrants from various European countries and African Americans from the American South contributing significantly to this growth. Over time, the population has fluctuated, with a decline in the latter half of the 20th century followed by a rise in recent years. \n",
            "\n",
            "As per the latest census estimates, the largest racial or ethnic groups in Chicago are non-Hispanic White, Black, and Hispanic. The city also has a significant LGBT population and is celebrated for its cultural diversity.\n",
            "========================\n"
          ]
        }
      ],
      "source": [
        "response = top_agent.query(\n",
        "    \"Tell the demographics of Houston, and then compare that with the\"\n",
        "    \" demographics of Chicago\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cfe1dd4c-8bfd-43d0-99bc-ca60861dc418",
      "metadata": {
        "id": "cfe1dd4c-8bfd-43d0-99bc-ca60861dc418",
        "outputId": "ab0dda35-21bc-4366-d353-2bcd255ba145"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Houston has a diverse population with a demographic makeup that includes non-Hispanic whites (23.3%), Hispanics and Latino Americans (45.8%), Blacks or African Americans (22.4%), and Asian Americans (6.5%). The largest Hispanic or Latino American ethnic group in Houston is Mexican Americans. Houston is also home to the largest African American community west of the Mississippi River and has a growing Muslim population.\n",
            "\n",
            "On the other hand, Chicago is also known for its diverse demographics. The city has a significant non-Hispanic White population, along with a substantial Black population and Hispanic population. Chicago is celebrated for its cultural diversity and has a significant LGBT population.\n",
            "\n",
            "Both Houston and Chicago have diverse populations, with a mix of different racial and ethnic groups contributing to their vibrant communities.\n"
          ]
        }
      ],
      "source": [
        "print(response)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "843c79cb-fc7a-4eac-ba7f-6d25c7a8f08e",
      "metadata": {
        "id": "843c79cb-fc7a-4eac-ba7f-6d25c7a8f08e",
        "outputId": "cea278a3-a807-47d9-d2c0-a2b20e3a7fe8"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Houston is the most populous city in Texas and the fourth-most populous city in the United States. It has a population of 2,304,580 as of the 2020 U.S. census. The city is known for its diversity, with a significant proportion of minorities. In 2019, non-Hispanic whites made up 23.3% of the population, Hispanics and Latino Americans 45.8%, Blacks or African Americans 22.4%, and Asian Americans 6.5%. The largest Hispanic or Latino American ethnic group in Houston is Mexican Americans, comprising 31.6% of the population.\n",
            "\n",
            "In comparison, Chicago is the third-most populous city in the United States. According to the 2020 U.S. census, Chicago has a population of 2,746,388. The demographics of Chicago are different from Houston, with non-Hispanic whites making up 32.7% of the population, Hispanics and Latino Americans 29.9%, Blacks or African Americans 29.8%, and Asian Americans 7.6%. The largest Hispanic or Latino American ethnic group in Chicago is Mexican Americans, comprising 21.6% of the population.\n",
            "\n",
            "Overall, both Houston and Chicago have diverse populations, but the specific demographic composition differs between the two cities.\n"
          ]
        }
      ],
      "source": [
        "# baseline\n",
        "response = base_query_engine.query(\n",
        "    \"Tell the demographics of Houston, and then compare that with the\"\n",
        "    \" demographics of Chicago\"\n",
        ")\n",
        "print(str(response))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e6f99183-bdd9-4d49-8980-44919cf86c03",
      "metadata": {
        "id": "e6f99183-bdd9-4d49-8980-44919cf86c03"
      },
      "outputs": [],
      "source": [
        "# baseline: the response tells you nothing about Chicago...\n",
        "response.source_nodes[3].get_content()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5b7e0a4c-b7be-4797-9c82-1577693bd117",
      "metadata": {
        "id": "5b7e0a4c-b7be-4797-9c82-1577693bd117",
        "outputId": "c0e85055-4331-42ca-eb6b-21474d29db84"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "=== Calling Function ===\n",
            "Calling function: tool_Shanghai with args: {\n",
            "  \"input\": \"history\"\n",
            "}\n",
            "=== Calling Function ===\n",
            "Calling function: vector_tool with args: {\n",
            "  \"input\": \"history\"\n",
            "}\n",
            "Got output: Shanghai has a rich history that dates back to ancient times. However, in the context provided, the history of Shanghai is mainly discussed in relation to its modern development. After the war, Shanghai's economy experienced significant growth, with increased agricultural and industrial output. The city's administrative divisions were rearranged, and it became a center for radical leftism during the 1950s and 1960s. The Cultural Revolution had a severe impact on Shanghai's society, but the city maintained economic production with a positive growth rate. Shanghai also played a significant role in China's Third Front campaign and has been a major contributor of tax revenue to the central government. Economic reforms were initiated in Shanghai in 1990, leading to the development of the Pudong district and its classification as an Alpha+ city.\n",
            "========================\n",
            "Got output: Shanghai's history is rich and complex, dating back to ancient times. However, its modern development is particularly noteworthy. After the war, Shanghai experienced significant economic growth, with a boost in both agricultural and industrial output. The city's administrative divisions were restructured, and it became a hub for radical leftism during the 1950s and 1960s.\n",
            "\n",
            "The Cultural Revolution had a profound impact on Shanghai's society, but despite this, the city managed to maintain economic production with a positive growth rate. Shanghai also played a significant role in China's Third Front campaign and has been a major contributor of tax revenue to the central government.\n",
            "\n",
            "In 1990, economic reforms were initiated in Shanghai, leading to the development of the Pudong district. This has helped Shanghai to be classified as an Alpha+ city, indicating its influence on the global economic stage.\n",
            "========================\n",
            "=== Calling Function ===\n",
            "Calling function: tool_Beijing with args: {\n",
            "  \"input\": \"history\"\n",
            "}\n",
            "=== Calling Function ===\n",
            "Calling function: vector_tool with args: {\n",
            "  \"input\": \"history\"\n",
            "}\n",
            "Got output: Beijing has a rich history that spans several dynasties. It was the capital of the Ming dynasty, during which the city took its current shape and many of its major attractions, such as the Forbidden City and the Temple of Heaven, were constructed. The Qing dynasty succeeded the Ming dynasty and made Beijing its sole capital. During this time, the Imperial residence and the general layout of the city remained largely unchanged. However, the city faced challenges during the Second Opium War and the Boxer Rebellion, resulting in the looting and destruction of important structures. In the early 20th century, Beijing saw the signing of a peace agreement between the Eight-Nation Alliance and the Chinese government, which led to the restoration of Qing dynasty rule. However, the dynasty eventually collapsed in 1911.\n",
            "========================\n",
            "Got output: Beijing has a rich and complex history that spans several dynasties. It served as the capital during the Ming dynasty, during which the city took its current shape and many of its major attractions, such as the Forbidden City and the Temple of Heaven, were constructed. The Qing dynasty succeeded the Ming dynasty and made Beijing its sole capital. During this time, the Imperial residence and the general layout of the city remained largely unchanged.\n",
            "\n",
            "However, the city faced significant challenges during the Second Opium War and the Boxer Rebellion, which resulted in the looting and destruction of important structures. In the early 20th century, Beijing saw the signing of a peace agreement between the Eight-Nation Alliance and the Chinese government, leading to the restoration of Qing dynasty rule. However, the dynasty eventually collapsed in 1911. Despite these tumultuous events, Beijing has managed to preserve its historical heritage while also evolving into a modern metropolis.\n",
            "========================\n",
            "=== Calling Function ===\n",
            "Calling function: tool_Shanghai with args: {\n",
            "  \"input\": \"current economy\"\n",
            "}\n",
            "=== Calling Function ===\n",
            "Calling function: vector_tool with args: {\n",
            "  \"input\": \"current economy\"\n",
            "}\n",
            "Got output: The current economy of Shanghai is strong and thriving. It is a global center for finance and innovation, and a national center for commerce, trade, and transportation. The city has a diverse economy, with its six largest industries comprising about half of its GDP. Shanghai has experienced rapid development and has been one of the fastest-developing cities in the world. It has recorded double-digit GDP growth in almost every year between 1992 and 2008. As of 2021, Shanghai had a GDP of CN¥4.46 trillion ($1.106 trillion in PPP), making it one of the wealthiest cities in China. It is also the most expensive city in mainland China to live in. Shanghai is a major player in the global financial industry, ranking first in Asia and third globally in the Global Financial Centres Index. It is home to the Shanghai Stock Exchange, the largest stock exchange in China and the fourth-largest in the world. The city has attracted significant foreign investment and has been a hub for the technology industry and startups. Overall, the current economy of Shanghai is robust and continues to grow.\n",
            "========================\n",
            "Got output: The current economy of Shanghai is robust and thriving. It is a global center for finance and innovation, and a national center for commerce, trade, and transportation. The city has a diverse economy, with its six largest industries comprising about half of its GDP. \n",
            "\n",
            "Shanghai has experienced rapid development and has been one of the fastest-developing cities in the world. It has recorded double-digit GDP growth in almost every year between 1992 and 2008. As of 2021, Shanghai had a GDP of CN¥4.46 trillion ($1.106 trillion in PPP), making it one of the wealthiest cities in China. \n",
            "\n",
            "Shanghai is also the most expensive city in mainland China to live in. It is a major player in the global financial industry, ranking first in Asia and third globally in the Global Financial Centres Index. The city is home to the Shanghai Stock Exchange, the largest stock exchange in China and the fourth-largest in the world. \n",
            "\n",
            "The city has attracted significant foreign investment and has been a hub for the technology industry and startups. Overall, the current economy of Shanghai is robust and continues to grow.\n",
            "========================\n",
            "=== Calling Function ===\n",
            "Calling function: tool_Beijing with args: {\n",
            "  \"input\": \"current economy\"\n",
            "}\n",
            "=== Calling Function ===\n",
            "Calling function: vector_tool with args: {\n",
            "  \"input\": \"current economy\"\n",
            "}\n",
            "Got output: The current economy of Beijing is dominated by the tertiary sector, which includes services such as professional services, wholesale and retail, information technology, commercial real estate, scientific research, and residential real estate. This sector generated 83.8% of the city's output in 2022. The secondary sector, which includes manufacturing and construction, accounted for 15.8% of output, while the primary sector, which includes agriculture and mining, contributed only 0.26%. The city has also identified six high-end economic output zones that are driving local economic growth, including Zhongguancun, Beijing Financial Street, Beijing Central Business District (CBD), Beijing Economic and Technological Development Area (Yizhuang), Beijing Airport Economic Zone, and Beijing Olympic Center Zone. These zones are home to various industries and sectors, such as technology companies, financial institutions, office buildings, industrial parks, and entertainment and sports centers.\n",
            "========================\n",
            "Got output: The current economy of Beijing is primarily driven by the tertiary sector, which includes services such as professional services, wholesale and retail, information technology, commercial real estate, scientific research, and residential real estate. This sector generated 83.8% of the city's output in 2022. The secondary sector, which includes manufacturing and construction, accounted for 15.8% of output, while the primary sector, which includes agriculture and mining, contributed only 0.26%.\n",
            "\n",
            "Beijing has also identified six high-end economic output zones that are driving local economic growth. These include Zhongguancun, Beijing Financial Street, Beijing Central Business District (CBD), Beijing Economic and Technological Development Area (Yizhuang), Beijing Airport Economic Zone, and Beijing Olympic Center Zone. These zones are home to various industries and sectors, such as technology companies, financial institutions, office buildings, industrial parks, and entertainment and sports centers.\n",
            "========================\n"
          ]
        }
      ],
      "source": [
        "response = top_agent.query(\n",
        "    \"Tell me the differences between Shanghai and Beijing in terms of history\"\n",
        "    \" and current economy\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "98540ec5-093d-41a0-888b-246ee7093cff",
      "metadata": {
        "id": "98540ec5-093d-41a0-888b-246ee7093cff",
        "outputId": "6712d84e-d825-417a-93ba-b4fe72763513"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "In terms of history, both Shanghai and Beijing have rich and complex pasts. Shanghai's history dates back to ancient times, but its modern development is particularly noteworthy. It experienced significant economic growth after the war and played a major role in China's economic reforms. Beijing, on the other hand, has a history that spans several dynasties and served as the capital during the Ming and Qing dynasties. It has preserved its historical heritage while evolving into a modern metropolis.\n",
            "\n",
            "In terms of current economy, Shanghai is a global center for finance and innovation. It has a diverse economy and has experienced rapid development, with a high GDP and significant foreign investment. It is a major player in the global financial industry and is home to the Shanghai Stock Exchange. Beijing's economy is primarily driven by the tertiary sector, with a focus on services such as professional services, information technology, and commercial real estate. It has identified high-end economic output zones that are driving local economic growth.\n",
            "\n",
            "Overall, both cities have thriving economies, but Shanghai has a stronger focus on finance and global influence, while Beijing has a diverse economy with a focus on services and high-end economic zones.\n"
          ]
        }
      ],
      "source": [
        "print(str(response))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a0f49e28-e85e-4c93-8eed-68122ad9f537",
      "metadata": {
        "id": "a0f49e28-e85e-4c93-8eed-68122ad9f537",
        "outputId": "b283112a-4ef2-4edb-e64d-578a494047eb"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Shanghai and Beijing have distinct differences in terms of history and current economy. Historically, Shanghai was the largest and most prosperous city in East Asia during the 1930s, while Beijing served as the capital of the Republic of China and later the People's Republic of China. Shanghai experienced significant growth and redevelopment in the 1990s, while Beijing expanded its urban area and underwent rapid development in the last two decades.\n",
            "\n",
            "In terms of the current economy, Shanghai is considered the \"showpiece\" of China's booming economy. It is a global center for finance and innovation, with a strong focus on industries such as retail, finance, IT, real estate, machine manufacturing, and automotive manufacturing. Shanghai is also home to the world's busiest container port, the Port of Shanghai. The city has a high GDP and is classified as an Alpha+ city by the Globalization and World Cities Research Network.\n",
            "\n",
            "On the other hand, Beijing is a global financial center and ranks third globally in the Global Financial Centres Index. It is also a hub for the Chinese and global technology industry, with a large startup ecosystem. Beijing has a strong presence in industries such as finance, technology, and pharmaceuticals. The city is home to the headquarters of large state banks and insurance companies, as well as the country's financial regulatory agencies.\n",
            "\n",
            "Overall, while both Shanghai and Beijing are important economic centers in China, Shanghai has a stronger focus on industries such as finance, retail, and manufacturing, while Beijing has a strong presence in finance, technology, and pharmaceuticals.\n"
          ]
        }
      ],
      "source": [
        "# baseline\n",
        "response = base_query_engine.query(\n",
        "    \"Tell me the differences between Shanghai and Beijing in terms of history\"\n",
        "    \" and current economy\"\n",
        ")\n",
        "print(str(response))"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "llama_index_v2",
      "language": "python",
      "name": "llama_index_v2"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}